Refactor streaming scan construction into `StreamingScan` static methods by Matt711 · Pull Request #22889 · rapidsai/cudf

Matt711 · 2026-06-15T12:27:25Z

Description

Small refactor that makes the logic easier to follow IMO. There's one StreamingScan per IR::Scan in the logical plan. Each of which has a sequence of either SplitScans or FusedScans.

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

coderabbitai · 2026-06-15T12:28:48Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 35830cdd-1d42-48e2-82d1-56f3cc14e016

📥 Commits

Reviewing files that changed from the base of the PR and between 7d00ac3 and 096ffcc.

📒 Files selected for processing (1)

python/cudf_polars/cudf_polars/streaming/io.py

📝 Walkthrough

Summary by CodeRabbit

Refactor
- Improved type safety and consistency in streaming file I/O handling through updated type annotations.
- Restructured internal scan expansion logic to streamline partition processing workflow and reduce code duplication.
Tests
- Updated streaming tests to verify refactored scan expansion behavior with explicit partition count calculations.

Walkthrough

expand_scan_for_rank now accepts an explicit partition_count parameter and returns a StreamingScan directly instead of a list. A new _rank_slice helper computes per-rank partition offsets. StreamingScan gains for_split_files and for_fused_files factory classmethods, changes its scans storage to Sequence, and updates do_evaluate's type annotation accordingly. Actor graph and test call sites are updated to match.

Changes

Rank-aware StreamingScan refactor

Layer / File(s)	Summary
`_rank_slice` helper, `expand_scan_for_rank`, and `StreamingScan` factory methods `python/cudf_polars/cudf_polars/streaming/io.py`	Adds `Self` and `Sequence` imports; introduces `_rank_slice(total, rank, nranks)` helper; changes `expand_scan_for_rank` to accept `partition_count` and return `StreamingScan` via dispatch to new `StreamingScan.for_split_files` / `for_fused_files` classmethods; updates `StreamingScan.__init__` to store `scans` as `Sequence[SplitScan] \| Sequence[FusedScan]`; updates `do_evaluate` annotation.
Scan lowering call site update `python/cudf_polars/cudf_polars/streaming/io.py`	Updates the `Scan` lowering path to pass `count` as `partition_count` and receive `new_ir` as a `StreamingScan` directly, removing the prior `StreamingScan(scans, ir)` wrapping step.
Actor graph typing updates `python/cudf_polars/cudf_polars/streaming/actor_graph/io.py`	Adds `Sequence` under `TYPE_CHECKING`, annotates `scans` with the union-of-sequences type in `scan_node`, and adds `type: ignore[arg-type]` to the `producer_tasks` append call.
Test updates for new `expand_scan_for_rank` contract `python/cudf_polars/tests/streaming/test_scan.py`	Adds `math` import; updates fused and split-files tests to compute and pass explicit `partition_count` and iterate over `streaming_scan.scans` in assertions.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

rapidsai/cudf#22758: Refactors streaming Scan lowering and the scan_node actor to consume StreamingScan, which this PR extends by redesigning expand_scan_for_rank and StreamingScan construction.
rapidsai/cudf#22838: Modifies FusedScan/SplitScan handling and StreamingScan's stored scans typing, directly related to the signature and type changes in this PR.

Suggested reviewers

TomAugspurger

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 72.73% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main refactor: consolidating streaming scan construction logic into StreamingScan static methods, which is the core change across the modified files.
Description check	✅ Passed	The description is related to the changeset, explaining the refactor's purpose and architectural pattern, though it lacks some technical detail about the specific changes.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

TomAugspurger

There's one StreamingScan per IR::Scan in the logical plan

Is this new on this PR, or also the behavior on main?

Matt711 · 2026-06-15T15:05:06Z

There's one StreamingScan per IR::Scan in the logical plan

Is this new on this PR, or also the behavior on main?

This is behavior on main :)

Co-authored-by: Tom Augspurger <tom.augspurger88@gmail.com>

coderabbitai

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@python/cudf_polars/cudf_polars/streaming/io.py`:
- Line 593: The `Self` type hint is used in the `for_split_files` method's
return type annotation but is not imported from `typing`, causing a `NameError`
at module import time. Add `Self` to the typing imports at the top of the file
(line 14). Additionally, there is an inconsistency in return type hints: the
`for_split_files` method uses `-> Self:` while the `for_fused_files` method uses
`-> StreamingScan:`. Update the `for_fused_files` method's return type
annotation to use `-> Self:` instead of `-> StreamingScan:` to maintain
consistency across both methods.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 36ba7600-9a6a-4f6e-826b-b74b73fd97a8

📥 Commits

Reviewing files that changed from the base of the PR and between 375728a and 7d00ac3.

📒 Files selected for processing (1)

python/cudf_polars/cudf_polars/streaming/io.py

coderabbitai · 2026-06-15T16:14:54Z

Note

Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it.

An unexpected error occurred while generating fixes: Not Found - https://docs.github.com/rest/git/refs#get-a-reference

Matt711 · 2026-06-15T16:48:06Z

/merge

Refactor Scan expansion in cudf-polars

fbd6df3

Matt711 added the improvement Improvement / enhancement to an existing function label Jun 15, 2026

Matt711 requested a review from a team as a code owner June 15, 2026 12:27

Matt711 added the non-breaking Non-breaking change label Jun 15, 2026

Matt711 requested a review from mroeschke June 15, 2026 12:27

Merge branch 'main' into imp/polars/refactor-expand-scan-for-rank

375728a

github-actions Bot assigned Matt711 Jun 15, 2026

github-actions Bot added Python Affects Python cuDF API. cudf-polars Issues specific to cudf-polars labels Jun 15, 2026

github-project-automation Bot added this to cuDF Python Jun 15, 2026

GPUtester moved this to In Progress in cuDF Python Jun 15, 2026

TomAugspurger approved these changes Jun 15, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/io.py Outdated

Update python/cudf_polars/cudf_polars/streaming/io.py

7d00ac3

Co-authored-by: Tom Augspurger <tom.augspurger88@gmail.com>

coderabbitai Bot reviewed Jun 15, 2026

View reviewed changes

Comment thread python/cudf_polars/cudf_polars/streaming/io.py

Matt711 and others added 2 commits June 15, 2026 16:47

use Sefl from typing

afa3557

Merge branch 'main' into imp/polars/refactor-expand-scan-for-rank

096ffcc

rapids-bot Bot merged commit 905dd3c into rapidsai:main Jun 15, 2026
106 checks passed

github-project-automation Bot moved this from In Progress to Done in cuDF Python Jun 15, 2026

Matt711 deleted the imp/polars/refactor-expand-scan-for-rank branch June 16, 2026 14:43

coderabbitai Bot mentioned this pull request Jul 9, 2026

Fetch parquet metadata in rapidsmpf streaming network #23145

Open

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Refactor streaming scan construction into `StreamingScan` static methods - #22889

Refactor streaming scan construction into `StreamingScan` static methods#22889
rapids-bot[bot] merged 5 commits into
rapidsai:mainfrom
Matt711:imp/polars/refactor-expand-scan-for-rank

Matt711 commented Jun 15, 2026

Uh oh!

coderabbitai Bot commented Jun 15, 2026 •

edited

Loading

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

TomAugspurger left a comment

Uh oh!

Uh oh!

Matt711 commented Jun 15, 2026

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

coderabbitai Bot commented Jun 15, 2026 •

edited

Loading

Uh oh!

Matt711 commented Jun 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

Matt711 commented Jun 15, 2026

Description

Checklist

Uh oh!

coderabbitai Bot commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

TomAugspurger left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Matt711 commented Jun 15, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Matt711 commented Jun 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coderabbitai Bot commented Jun 15, 2026 •

edited

Loading

coderabbitai Bot commented Jun 15, 2026 •

edited

Loading